import os
import re

import numpy as np
import tiktoken
from tqdm import trange,tqdm
import time
from src.model import APIModel
from src.database import database
from src.utils import tokenCounter, max_retry_times
from src.prompt import ROUGH_OUTLINE_PROMPT, MERGING_OUTLINE_PROMPT, SUBSECTION_OUTLINE_PROMPT, EDIT_FINAL_OUTLINE_PROMPT, QUERY_GENERATE_PROMPT, BOOL_QUERY_GENERATE_PROMPT

class outlineWriter():
    
    def __init__(self, model:str, api_key:str, database) -> None:
        
        self.model, self.api_key = model, api_key
        self.api_model = APIModel(model=self.model, api_key=self.api_key)

        self.db = database
        self.token_counter = tokenCounter()
        self.input_token_usage, self.output_token_usage = 0, 0


    def draft_outline(self, topic, reference_num=600, chunk_size=30000, section_num=6, rag_num=20):
        # Get database
        # 生成问题，然后从问题进行索引
        #changed
        bool_queries, nl_queries = self.query_generate(topic)  #change
        references1 = self.db.get_papers_from_queries(topic, bool_queries, num=1000)   #change
        references2 = self.db.get_chunks_from_queries(nl_queries, num=reference_num)
        references = self.db.merge_references(references1, references2)
        references = self.db.rerank_ref(references)[:reference_num]
        references_contents = [r['chunk_text'] for r in references]
        references_titles = [r['paper_title'] for r in references]
        paper_chunks, titles_chunks = self.chunking(references_contents, references_titles, chunk_size=chunk_size)

        # generate rough section-level outline
        outlines = self.generate_rough_outlines(topic=topic, papers_chunks=paper_chunks, titles_chunks=titles_chunks, section_num=section_num)
        
        # merge outline
        section_outline = self.merge_outlines(topic=topic, outlines=outlines, section_num=section_num)

        # generate subsection-level outline
        subsection_outlines = self.generate_subsection_outlines(topic=topic, section_outline= section_outline,rag_num= rag_num)
        
        merged_outline = self.process_outlines(section_outline, subsection_outlines)
        
        # edit final outline
        final_outline = self.edit_final_outline(merged_outline)

        return final_outline
    
    # 这个方法貌似没有用到
    # def without_merging(self, topic, reference_num = 600, chunk_size = 30000, section_num = 6):
    #     references_ids = self.db.get_ids_from_topic(topic = topic, num = reference_num, shuffle = False)
    #     references_infos = self.db.get_paper_info_from_ids(references_ids)

    #     references_titles = [r['title'] for r in references_infos]
    #     references_papers = [r['abs'] for r in references_infos]
    #     papers_chunks, titles_chunks = self.chunking(references_papers, references_titles, chunk_size=chunk_size)

    #     # generate rough section-level outline
    #     section_outline = self.generate_rough_outlines(topic=topic, papers_chunks = [papers_chunks[0]], titles_chunks = [titles_chunks[0]], section_num=section_num)[0]
        
    #     # generate subsection-level outline
    #     subsection_outlines = self.generate_subsection_outlines(topic=topic, section_outline= section_outline)
        
    #     final_outline = self.process_outlines(section_outline, subsection_outlines)

    #     return final_outline, section_outline, subsection_outlines

    def compute_price(self):
        return self.token_counter.compute_price(input_tokens=self.input_token_usage, output_tokens=self.output_token_usage, model=self.model)

    def generate_rough_outlines(self, topic, papers_chunks, titles_chunks, section_num = 8):
        '''
        You want to write an overall and comprehensive academic survey about "[TOPIC]".\n
        You are provided with a list of papers related to the topic below:\n
        ---
        [PAPER LIST]
        ---
        You need to draft an outline based on the given papers.
        The outline should contain a title and several sections.
        Each section follows with a brief sentence to describe what to write in this section.
        The outline is supposed to be comprehensive and contains [SECTION NUM] sections.

        Return in the format:
        <format>
        Title: [TITLE OF THE SURVEY]
        Section 1: [NAME OF SECTION 1]
        Description 1: [DESCRIPTION OF SECTION 1]

        Section 2: [NAME OF SECTION 2]
        Description 2: [DESCRIPTION OF SECTION 2]

        ...

        Section [SECTION NUM]: [NAME OF SECTION K]
        Description [SECTION NUM]: [DESCRIPTION OF SECTION K]
        </format>
        The outline:
        '''

        prompts = []

        for i in trange(len(papers_chunks)):
            titles = titles_chunks[i]
            papers = papers_chunks[i]
            paper_texts = '' 
            for i, t, p in zip(range(len(papers)), titles, papers):
                paper_texts += f'---\npaper_title: {t}\n\npaper_content:\n\n{p}\n'
            paper_texts+='---\n'

            prompt = self.__generate_prompt(ROUGH_OUTLINE_PROMPT, paras={'PAPER LIST': paper_texts, 'TOPIC': topic, 'SECTION NUM': str(section_num)})
            prompts.append(prompt)
        self.input_token_usage += self.token_counter.num_tokens_from_list_string(prompts)
        outlines = self.api_model.batch_chat(text_batch=prompts, temperature=1)
        self.output_token_usage += self.token_counter.num_tokens_from_list_string(outlines)
        return outlines
    
    def merge_outlines(self, topic, outlines, section_num):
        #print(section_num)
        '''
        You are an expert in artificial intelligence who wants to write an overall survey about [TOPIC].\n
        You are provided with a list of outlines as candidates below:\n
        ---
        [OUTLINE LIST]
        ---
        Each outline contains a title and several sections.\n
        Each section follows with a brief sentence to describe what to write in this section.\n\n
        You need to generate a final outline with [SECTION NUM] sections based on these provided outlines.\n
        Return in the format:
        <format>
        Title: [TITLE OF THE SURVEY]
        Section 1: [NAME OF SECTION 1]
        Description 1: [DESCRIPTION OF SECTION 1]

        Section 2: [NAME OF SECTION 2]
        Description 2: [DESCRIPTION OF SECTION 2]

        ...

        Section [SECTION NUM]: [NAME OF SECTION K]
        Description [SECTION NUM]: [DESCRIPTION OF SECTION K]
        </format>
        Only return the final outline without any other information:
        '''
        outline_texts = '' 
        for i, o in zip(range(len(outlines)), outlines):
            outline_texts += f'---\noutline_id: {i}\n\noutline_content:\n\n{o}\n'
        outline_texts+='---\n'
        prompt = self.__generate_prompt(MERGING_OUTLINE_PROMPT, paras={'OUTLINE LIST': outline_texts, 'TOPIC':topic, 'SECTION NUM': str(section_num)})
        self.input_token_usage += self.token_counter.num_tokens_from_string(prompt)

        outline = self.api_model.chat(prompt, temperature=1)
        self.output_token_usage += self.token_counter.num_tokens_from_string(outline)
        return outline
    
    def generate_subsection_outlines(self, topic, section_outline, rag_num):
        '''
        You are an expert in artificial intelligence who wants to write an overall survey about [TOPIC].
        You have created an overall outline below:
        ---
        [OVERALL OUTLINE]
        ---
        The outline contains a title and several sections.
        Each section follows with a brief sentence to describe what to write in this section.
        <instruction>
        You need to enrich the section [SECTION NAME].
        The description of [SECTION NAME]: [SECTION DESCRIPTION]
        You need to generate the framework containing several subsections based on the overall outlines.
        Each subsection follows with a brief sentence to describe what to write in this subsection.
        These papers provided for references:
        ---
        [PAPER LIST]
        ---
        Note that the number of subsections is no more than 4.
        Return the outline in the format:
        <format>
        Subsection 1: [NAME OF SUBSECTION 1]
        Description 1: [DESCRIPTION OF SUBSECTION 1]

        Subsection 2: [NAME OF SUBSECTION 2]
        Description 2: [DESCRIPTION OF SUBSECTION 2]

        ...

        Subsection K: [NAME OF SUBSECTION K]
        Description K: [DESCRIPTION OF SUBSECTION K]
        </format>
        </instruction>
        Only return the outline without any other information:
        '''


        survey_title, survey_sections, survey_section_descriptions = self.extract_title_sections_descriptions(section_outline)

        prompts = []

        for section_name, section_description in zip(survey_sections, survey_section_descriptions):
            #changed
            bool_queries, nl_queries = self.query_generate(section_description)  # change
            references1 = self.db.get_papers_from_queries(section_description, bool_queries, num=1000)  # change
            references2 = self.db.get_chunks_from_queries(nl_queries, num=rag_num)  # change
            references = self.db.merge_references(references1, references2)
            references = self.db.rerank_ref(references)[:rag_num]
            references_papers = [r['chunk_text'] for r in references]
            references_titles = [r['paper_title'] for r in references]

            paper_texts = '' 
            for i, t, p in zip(range(len(references_papers)), references_titles, references_papers):
                paper_texts += f'---\npaper_title: {t}\n\npaper_content:\n\n{p}\n'
            paper_texts+='---\n'
            prompt = self.__generate_prompt(SUBSECTION_OUTLINE_PROMPT, paras={'OVERALL OUTLINE': section_outline,'SECTION NAME': section_name,
                                                                          'SECTION DESCRIPTION':section_description,'TOPIC':topic,'PAPER LIST':paper_texts})
            prompts.append(prompt)
        self.input_token_usage += self.token_counter.num_tokens_from_list_string(prompts)

        sub_outlines = self.api_model.batch_chat(prompts, temperature=1)

        self.output_token_usage += self.token_counter.num_tokens_from_list_string(sub_outlines)
        return sub_outlines

    def edit_final_outline(self, outline):
        '''
        You are an expert in artificial intelligence who wants to write a overall survey about [TOPIC].\n
        You have created a draft outline below:\n
        ---
        [OVERALL OUTLINE]
        ---
        The outline contains a title and several sections.\n
        Each section follows with a brief sentence to describe what to write in this section.\n\n
        Under each section, there are several subsections.
        Each subsection also follows with a brief sentence of description.
        You need to modify the outline to make it both comprehensive and coherent with no repeated subsections.
        Return the final outline in the format:
        <format>
        # [TITLE OF SURVEY]

        ## [NAME OF SECTION 1]

        ### [NAME OF SUBSECTION 1]

        ### [NAME OF SUBSECTION 2]

        ...

        ### [NAME OF SUBSECTION L]
        
        ## [NAME OF SECTION 2]

        ...

        ## [NAME OF SECTION K]

        ...
        </format>
        Only return the final outline without any other information:
        '''

        prompt = self.__generate_prompt(EDIT_FINAL_OUTLINE_PROMPT, paras={'OVERALL OUTLINE': outline})
        self.input_token_usage += self.token_counter.num_tokens_from_string(prompt)
        outline = self.api_model.chat(prompt, temperature=1)
        self.output_token_usage += self.token_counter.num_tokens_from_string(outline)
        return outline.replace('<format>\n','').replace('</format>','')
 
    def __generate_prompt(self, template, paras):
        prompt = template
        for k in paras.keys():
            prompt = prompt.replace(f'[{k}]', paras[k])
        return prompt
    
    def extract_title_sections_descriptions(self, outline):
        title = outline.split('Title: ')[1].split('\n')[0]
        sections, descriptions = [], []
        for i in range(100):
            if f'Section {i+1}' in outline:
                sections.append(outline.split(f'Section {i+1}: ')[1].split('\n')[0])
                descriptions.append(outline.split(f'Description {i+1}: ')[1].split('\n')[0])
        return title, sections, descriptions
    
    def extract_subsections_subdescriptions(self, outline):
        subsections, subdescriptions = [], []
        for i in range(100):
            if f'Subsection {i+1}' in outline:
                subsections.append(outline.split(f'Subsection {i+1}: ')[1].split('\n')[0])
                subdescriptions.append(outline.split(f'Description {i+1}: ')[1].split('\n')[0])
        return subsections, subdescriptions
    
    def chunking(self, papers, titles, chunk_size = 14000):
        paper_chunks, title_chunks = [], []
        total_length = self.token_counter.num_tokens_from_list_string(papers)
        num_of_chunks = int(total_length / chunk_size) + 1
        avg_len = int(total_length / num_of_chunks) + 1
        split_points = []
        l = 0
        for j in range(len(papers)):
            l += self.token_counter.num_tokens_from_string(papers[j])
            if l > avg_len:
                l = 0
                split_points.append(j)
                continue
        start = 0
        for point in split_points:
            paper_chunks.append(papers[start:point])
            title_chunks.append(titles[start:point])
            start = point
        paper_chunks.append(papers[start:])
        title_chunks.append(titles[start:])
        return paper_chunks, title_chunks
       
    def process_outlines(self, section_outline, sub_outlines):
        res = ''
        survey_title, survey_sections, survey_section_descriptions = self.extract_title_sections_descriptions(outline=section_outline)
        res += f'# {survey_title}\n\n'
        for i in range(len(survey_sections)):
            section = survey_sections[i]
            res += f'## {i+1} {section}\nDescription: {survey_section_descriptions[i]}\n\n'
            subsections, subsection_descriptions = self.extract_subsections_subdescriptions(sub_outlines[i])
            for j in range(len(subsections)):
                subsection = subsections[j]
                res += f'### {i+1}.{j+1} {subsection}\nDescription: {subsection_descriptions[j]}\n\n'
        return res


    '''def query_generate(self, topic, q_num=5):
        prompt = self.__generate_prompt(QUERY_GENERATE_PROMPT, paras={'TOPIC': topic, 'n': str(q_num)})
        retry_times = 0
        while retry_times < max_retry_times:
            try:
                self.input_token_usage += self.token_counter.num_tokens_from_string(prompt)
                ans = self.api_model.chat(prompt, temperature=1)
                self.output_token_usage += self.token_counter.num_tokens_from_string(ans)
                # 提取 <begin> 和 <end> 之间的内容
                match = re.search(r"(?<=<begin>)[\s\S]*?(?=<end>)", ans, re.DOTALL)
                if match:
                    queries_text = match.group().strip()

                    # 提取并去除序号
                    queries = re.findall(r"\d+\.\s*([^0-9.][^.\n]*)", queries_text, re.DOTALL)
                    queries.append(topic)
                break
            except Exception as e:
                print(f'Error occurred during query generation: {e}')
                print(f'answer: {ans}')
                retry_times += 1
        if retry_times == max_retry_times:
            print('rewrite failed!')
            return [topic]
        return queries'''

    def query_generate(self, topic, q_num=5):
        prompt1 = self.__generate_prompt(BOOL_QUERY_GENERATE_PROMPT, paras={'TOPIC': topic})
        prompt2 = self.__generate_prompt(QUERY_GENERATE_PROMPT, paras={'TOPIC': topic, 'n': str(q_num)})
        all_queries = []
        for id, prompt in enumerate([prompt1, prompt2]):
            retry_times = 0
            while retry_times < max_retry_times:
                try:
                    self.input_token_usage += self.token_counter.num_tokens_from_string(prompt)
                    ans = self.api_model.chat(prompt, temperature=1)
                    self.output_token_usage += self.token_counter.num_tokens_from_string(ans)
                    # 提取 <begin> 和 <end> 之间的内容
                    match = re.search(r"(?<=<begin>)[\s\S]*?(?=<end>)", ans, re.DOTALL)
                    if match:
                        queries_text = match.group().strip()

                        # 提取并去除序号
                        queries = re.findall(r"\d+\.\s*([^0-9.][^.\n]*)", queries_text, re.DOTALL)

                        def remove_punctuation_except_hyphen(word):
                            # 使用正则表达式，保留连字符，去除其他标点符号
                            cleaned_word = re.sub(r"[^\w\s-]", "", word)
                            return cleaned_word
                        if id == 0:
                            queries = [remove_punctuation_except_hyphen(q) for q in queries]
                        if id == 1:
                            queries.append(topic)
                        queries = list(set(queries))
                        all_queries.append(queries)
                    break
                except Exception as e:
                    print(f'Error occurred during query generation: {e}')
                    print(f'answer: {ans}')
                    retry_times += 1
                if retry_times == max_retry_times:
                    print('rewrite failed!')
                    all_queries.append([topic])
        return all_queries[0], all_queries[1]   #bool_query, query